
#include <arpa/inet.h>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <ctype.h>
#include <fcntl.h>
#include <iostream>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

static int const ITERATIONS = 10000;
static int const MMSG_COUNT = 50;
static int const PORT = 5002;

int main(int argc, char const** argv)
{
    char buffer[2048];
    memset(buffer, 0, sizeof(buffer));

    // Arguments
    if (argc != 4)
    {
        std::cerr << "ARGUMENTS: <packet size> <packet count> <one|many>\n";
        return -1;
    }
    int const packet_size = atoi(argv[1]);
    int const max_packet_count = atoi(argv[2]);
    bool single_receive = true;
    if (std::string(argv[3]) == "one")
        single_receive = true;
    else if (std::string(argv[3]) == "many")
        single_receive = false;
    else
    {
        std::cerr << "Illegal one|many argument\n";
        return -1;
    }


    // Receive socket
    int receiver = socket(AF_INET, SOCK_DGRAM, 0);
    struct sockaddr_in receive_addr, remote_addr;
    memset(&receive_addr, 0, sizeof(receive_addr));
    receive_addr.sin_family = AF_INET;
    receive_addr.sin_port = htons(PORT);
    if (bind(receiver, (struct sockaddr*)&receive_addr, sizeof(receive_addr)))
    {
        std::cout << "bind error\n";
        return -1;
    }
    int flags = fcntl(receiver, F_GETFL, 0);
    fcntl(receiver, F_SETFL, flags | O_NONBLOCK);
    int recvbuf_size = 20000000;
    setsockopt(receiver, SOL_SOCKET, SO_RCVBUF, &recvbuf_size, sizeof(recvbuf_size));

    // Send socket
    int sender = socket(AF_INET, SOCK_DGRAM, 0);
    struct sockaddr_in send_addr;
    memset(&send_addr, 0, sizeof(send_addr));
    send_addr.sin_family = AF_INET;
    send_addr.sin_port = htons(PORT);
    inet_pton(AF_INET, "127.0.0.1", &send_addr.sin_addr);

    // recvmmsg setup
    struct mmsghdr msgs[MMSG_COUNT];
    struct iovec iovecs[MMSG_COUNT];
    char mmsg_bufs[MMSG_COUNT][2048];
    memset(msgs, 0, sizeof(msgs));
    for(int i = 0; i != MMSG_COUNT; ++i)
    {
        iovecs[i].iov_base = mmsg_bufs[i];
        iovecs[i].iov_len = 2048;
        msgs[i].msg_hdr.msg_iov = iovecs + i;
        msgs[i].msg_hdr.msg_iovlen = 1;
    }

    for(int packet_count = 1; packet_count <= max_packet_count; ++packet_count)
    {
    unsigned long long total_nanos = 0;
    for(int i = 0; i != ITERATIONS; ++i)
    {
        // Send packets
        for(int j = 0; j != packet_count; ++j)
        {
            if (sendto(sender, buffer, packet_size, 0, (struct sockaddr*)&send_addr, sizeof(send_addr)) != packet_size)
                std::cout << "sendto failure!\n";
        }

        // Receive packets
        struct timespec before, after;
        clock_gettime(CLOCK_MONOTONIC, &before);
        if (single_receive)
        {
            for(int j = 0; j < packet_count; )
                j += recvfrom(receiver, buffer, sizeof(buffer), 0, NULL, NULL) > 0 ? 1 : 0;
        }
        else
        {
            for(int j = 0; j < packet_count; )
            {
                j += recvmmsg(receiver, msgs, MMSG_COUNT, 0, NULL);
            }
        }
        clock_gettime(CLOCK_MONOTONIC, &after);
        long long nanoseconds = (after.tv_sec - before.tv_sec) * 1000000000 + (after.tv_nsec - before.tv_nsec);
        //std::cout << nanoseconds << std::endl;
        //std::cout.flush();
        total_nanos += nanoseconds;
    }
    std::cout << packet_count <<" packets: " << total_nanos / ITERATIONS << std::endl;
    }
}

